package constantReader;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.rmi.RemoteException;
import java.util.*;
import java.util.function.Function;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class ClassDependencyAnalyzer {

    private static final int CONSTANT_Utf8 = 1;
    private static final int CONSTANT_CLASS = 7;
    private static final Pattern CLASS_NAME_PATTERN = Pattern.compile("L([^;]+);");

    public static void main(String[] args) {
//        System.out.println(getClassesRelationMap(List.of(".")));
        System.out.println(getClassesRelationMap(List.of(".")));
    }

    public static Map<String, ClassRefs> getClassesRelationMap(List<String> pathList) {
        List<byte[]> classesList = getClassList(pathList);
        // key位类名，value为父类、继承接口
        return classesList
                .stream()
                .map(ClassDependencyAnalyzer::getClassNameParentAndInterface)
                .collect(Collectors.toMap(ClassRefs::getClsName, cls -> cls));
    }

    public static List<byte[]> getClassList(List<String> pathList) {
        List<byte[]> classesList = new ArrayList<>();
        // 文件处理函数
        Map<String, Function<File, List<byte[]>>> handlers = new HashMap<>() {
            {
                put("class", file -> {
                    try(InputStream inputStream = new FileInputStream(file)) {
                        return Collections.singletonList(inputStream.readAllBytes());
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                });
                put("jar", ClassDependencyAnalyzer::getClassListFromJar);
            }
        };
        // 提取需要遍历的文件,只有jar或者
        List<File> allFiles = dfsGetFiles(pathList, handlers.keySet().stream().toList());
        for (File file : allFiles) {
            classesList.addAll(handlers.get(getSuffix(file)).apply(file));
        }
        return classesList;
    }

    public static List<File> dfsGetFiles(List<String> pathList, List<String> suffixes) {
        List<File> files = new ArrayList<>();
        for (String path : pathList) {
            File file = new File(path);
            if (!file.exists()) continue;
            if (file.isFile() && (suffixes == null || suffixes.stream().anyMatch(suf -> suf.equals(getSuffix(file))))) files.add(file);
            if (file.isDirectory()) files.addAll(dfsGetFiles(Arrays.stream(Objects.requireNonNull(file.listFiles())).map(File::getAbsolutePath).toList(), suffixes));
        }
        return files;
    }

    public static List<byte[]> getClassListFromJar(File file) {
        return getClassesFromJar(file).values().stream().toList();
    }

    public static Map<String, byte[]> getClassesFromJar(File file) {
        Map<String, byte[]> clsMap = new HashMap<>();
        try (JarFile jarFile = new JarFile(file.getAbsoluteFile())) {
            Enumeration<JarEntry> entries = jarFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName();
                if (entryName.endsWith(".class")) {
                    clsMap.put(entryName.replaceAll("/", "."), jarFile.getInputStream(entry).readAllBytes());
                }
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return clsMap;
    }

    public static ClassRefs getClassNameParentAndInterface(byte[] bytes) {
        try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(bytes))) {
            // 读取魔数
            int magic = dis.readInt();
            if (magic != 0xCAFEBABE) {
                throw new IOException("不是有效的 .class 文件");
            }
            // 读取版本号
            dis.readUnsignedShort();
            dis.readUnsignedShort();
            // 读取常量池大小
            int constantPoolCount = dis.readUnsignedShort();
            Object[] constantPool = new Object[constantPoolCount];

            for (int i = 1; i < constantPoolCount; i++) {
                byte tag = dis.readByte();
                switch (tag) {
                    case CONSTANT_Utf8:
                        int length = dis.readUnsignedShort();
                        byte[] strBytes = new byte[length];
                        dis.readFully(strBytes);
                        constantPool[i] = new String(strBytes, StandardCharsets.UTF_8);
                        break;
                    case CONSTANT_CLASS:
                        constantPool[i] = dis.readShort();
                        break;
                    default:
                        i = skipConstant(dis, tag, i);
                }
            }
            // 读取访问标志
            dis.readUnsignedShort();

            // 读取类索引
            int thisClassIndex = dis.readUnsignedShort();
            short clsNameIndex = (short) constantPool[thisClassIndex];
            String className = ((String) constantPool[clsNameIndex]).replaceAll("/", ".");

            // 读取父类索引
            int superClassIndex = dis.readUnsignedShort();
            short superClassNameIndex = (short) constantPool[superClassIndex];
            String superClassName = ((String) constantPool[superClassNameIndex]).replaceAll("/", ".");

            // 读取接口数量
            int interfacesCount = dis.readUnsignedShort();
            List<String> interfaceNames = new ArrayList<>();
            for (int i = 0; i < interfacesCount; i++) {
                int interfaceIndex = dis.readUnsignedShort();
                short interfaceNameIndex = (short) constantPool[interfaceIndex];
                interfaceNames.add(((String) constantPool[interfaceNameIndex]).replaceAll("/", "."));
            }

            // 读取属性
            int fieldsCount = dis.readUnsignedShort();
            Set<String> fieldClsNames = new HashSet<>();
            for (int i = 0; i < fieldsCount; i++) {
                int fieldAccessFlags = dis.readUnsignedShort();
                int fieldNameIndex = dis.readUnsignedShort();
                int fieldDescriptorIndex = dis.readUnsignedShort();
                int attrLength = dis.readUnsignedShort();
                skipAttributes(dis, attrLength);
                String desc = (String) constantPool[fieldDescriptorIndex];
                if (desc.startsWith("L") && desc.endsWith(";")) {
                    fieldClsNames.add(desc.substring(1, desc.length()-1).replaceAll("/", "."));
                }
            }

            // 读取方法
            int methodsCount = dis.readUnsignedShort();
            Set<String> methodClsNames = new HashSet<>();
            for (int i = 0; i < methodsCount; i++) {
                int methodAccessFlags = dis.readUnsignedShort();
                int methodNameIndex = dis.readUnsignedShort();
                int methodDescriptorIndex = dis.readUnsignedShort();
                int attrLength = dis.readUnsignedShort();
                skipAttributes(dis, attrLength);
                String desc = (String) constantPool[methodDescriptorIndex];
                Matcher matcher = CLASS_NAME_PATTERN.matcher(desc);
                while(matcher.find()) {
                    methodClsNames.add(matcher.group(1).replaceAll("/", "."));
                }
            }

            return new ClassRefs(className, superClassName, interfaceNames, fieldClsNames, methodClsNames);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static int skipConstant(DataInputStream dis, byte tag, int i) throws IOException {
        switch (tag) {
            case 3: // CONSTANT_Integer
            case 4: // CONSTANT_Float
                dis.skipBytes(4);
                break;
            case 5: // CONSTANT_Long
            case 6: // CONSTANT_Double
                dis.skipBytes(8);
                i++;
                break;
            case 7: // CONSTANT_Class
            case 8: // CONSTANT_String
            case 16: // CONSTANT_MethodType
            case 19: // CONSTANT_Module
            case 20: // CONSTANT_Package
                dis.skipBytes(2);
                break;
            case 9: // CONSTANT_Fieldref
            case 10: // CONSTANT_Methodref
            case 11: // CONSTANT_InterfaceMethodref
            case 12: // CONSTANT_NameAndType
            case 18: // CONSTANT_InvokeDynamic
                dis.skipBytes(4);
                break;
            case 15: // CONSTANT_MethodHandle
                dis.skipBytes(3);
                break;
            case 21: // CONSTANT_UnresolvedClass
                dis.skipBytes(2);
                break;
            default:
                throw new RemoteException("未知标签");
        }
        return i;
    }

    private static void skipAttributes(DataInputStream dis, int attributesCount) throws IOException {
        for (int i = 0; i < attributesCount; i++) {
            // 读取属性名索引并跳过
            dis.readUnsignedShort();
            // 读取属性长度并跳过属性内容
            int attributeLength = dis.readInt();
            dis.skipBytes(attributeLength);
        }
    }

    private static String getSuffix(File file) {
        return getSuffix(file.getAbsolutePath());
    }

    private static String getSuffix(String fileName) {
        return fileName.substring(fileName.lastIndexOf(".")+1);
    }

    public static class Pair<L, R> {
        private L left;
        private R right;

        public Pair() {}

        public Pair (L left, R right) {
            this.left = left;
            this.right = right;
        }

        public L getLeft() {
            return left;
        }

        public void setLeft(L left) {
            this.left = left;
        }

        public R getRight() {
            return right;
        }

        public void setRight(R right) {
            this.right = right;
        }

        @Override
        public String toString() {
            return "Pair{" +
                    "left=" + left +
                    ", right=" + right +
                    '}';
        }

        public static <L, R> Pair<L, R> of(L left, R right) {
            return new Pair<>(left, right);
        }
    }

    public static class ClassRefs {
        private String clsName;
        private String parentName;
        private List<String> interfaceNames;
        private Set<String> fieldClassNames;
        private Set<String> methodClassNames;

        public ClassRefs(String clsName, String parentName, List<String> interfaceNames, Set<String> fieldClassNames, Set<String> methodClassNames) {
            this.clsName = clsName;
            this.parentName = parentName;
            this.interfaceNames = interfaceNames;
            this.fieldClassNames = fieldClassNames;
            this.methodClassNames = methodClassNames;
        }

        public String getClsName() {
            return clsName;
        }

        public void setClsName(String clsName) {
            this.clsName = clsName;
        }

        public String getParentName() {
            return parentName;
        }

        public void setParentName(String parentName) {
            this.parentName = parentName;
        }

        public List<String> getInterfaceNames() {
            return interfaceNames;
        }

        public void setInterfaceNames(List<String> interfaceNames) {
            this.interfaceNames = interfaceNames;
        }

        public Set<String> getFieldClassNames() {
            return fieldClassNames;
        }

        public void setFieldClassNames(Set<String> fieldClassNames) {
            this.fieldClassNames = fieldClassNames;
        }

        public Set<String> getMethodClassNames() {
            return methodClassNames;
        }

        public void setMethodClassNames(Set<String> methodClassNames) {
            this.methodClassNames = methodClassNames;
        }

        @Override
        public String toString() {
            return "ClassRefs{" +
                    "clsName='" + clsName + '\'' +
                    ", parentName='" + parentName + '\'' +
                    ", interfaceNames=" + interfaceNames +
                    ", fieldClassNames=" + fieldClassNames +
                    ", methodClassNames=" + methodClassNames +
                    '}';
        }
    }
}